import os
import cv2
from datetime import datetime
from ultralytics import YOLO
import matplotlib.pyplot as plt


def picture_flag(model, picture_path, result_image_dir):
    """
    使用YOLO模型标注图片并保存结果

    参数:
        model: 训练好的YOLO模型
        picture_path: 需要处理的图片完整路径
        result_image_dir: 结果保存目录

    返回:
        RGB格式的图像数据
    """
    # 确保输出目录存在
    os.makedirs(result_image_dir, exist_ok=True)

    # 加载图片
    img = cv2.imread(picture_path)
    if img is None:
        raise ValueError(f"无法加载图片: {picture_path}")

    # 使用模型进行预测
    results = model.predict(img)

    # 在图像上绘制标注结果
    annotated_img = results[0].plot()  # 获取带标注的图像

    # 生成输出文件名
    current_time = datetime.now().strftime("%m%d%H%M%S")
    output_filename = f"{current_time}_0.jpg"
    output_path = os.path.join(result_image_dir, output_filename)

    # 保存标注后的图像
    cv2.imwrite(output_path, annotated_img)

    # 将BGR转换为RGB
    rgb_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)

    return rgb_img

# def display_image(image):
#     """使用matplotlib展示图像"""
#     plt.figure(figsize=(10, 10))
#     plt.imshow(image)
#     plt.axis('off')  # 不显示坐标轴
#     plt.show()

def main():
    # 加载训练好的模型
    model = YOLO("best.pt")

    # 图片路径
    image_path = "*"

    # 结果保存目录
    result_dir = os.path.join("picture_inference", "picture_flag")

    # 进行检测
    result_image = picture_flag(model, image_path, result_dir)

    # 展示结果
    # display_image(result_image)

    print(f"检测完成，结果已保存到: {result_dir}")


if __name__ == "__main__":
    main()